import cv2
import os
import shutil
from pathlib import Path

class YOLOv5LabelingTool:
    def __init__(self):
        self.labels = {}  # 标签ID到标签名的映射
        self.current_image = None
        self.current_image_path = ""
        self.image_folder_path = ""  # 图片文件夹路径
        self.image_save_path = ""
        self.label_save_path = ""
        self.current_boxes = []  # 当前图片的所有标注框
        self.temp_point = None  # 临时存储第一个点击点
        self.current_label_id = None
        self.window_name = "YOLOv5 标注工具"
        self.image_files = []  # 按数字排序的图片文件列表
        self.current_image_index = 0  # 当前图片索引
        self.start_image_number = 1  # 开始标注的图片编号
        
    def setup_labels(self):
        """设置标签ID和名称"""
        print("=" * 60)
        print("YOLOv5 标注工具 - 标签设置")
        print("=" * 60)
        print("请输入标签信息 (输入 'done' 完成设置)")
        print("格式: ID号 标签名称")
        print("例如: 0 person")
        print("-" * 60)
        
        while True:
            user_input = input("请输入标签 (ID 名称): ").strip()
            if user_input.lower() == 'done':
                break
                
            try:
                parts = user_input.split(' ', 1)
                if len(parts) != 2:
                    print("❌ 格式错误，请使用: ID 名称")
                    continue
                    
                label_id = int(parts[0])
                label_name = parts[1]
                self.labels[label_id] = label_name
                print(f"✅ 已添加标签: {label_id} -> {label_name}")
                
            except ValueError:
                print("❌ ID必须是数字")
                continue
        
        if not self.labels:
            print("❌ 未设置任何标签")
            return False
            
        print("\n当前标签列表:")
        for label_id, label_name in sorted(self.labels.items()):
            print(f"  {label_id}: {label_name}")
        print("-" * 60)
        return True
    
    def setup_paths(self):
        """设置图片文件夹和标签保存路径"""
        print("\n路径设置:")

        # 图片文件夹路径
        while True:
            self.image_folder_path = input("请输入图片文件夹路径: ").strip()
            if not os.path.exists(self.image_folder_path):
                print("❌ 文件夹不存在")
                continue
            if not os.path.isdir(self.image_folder_path):
                print("❌ 这不是一个文件夹")
                continue

            # 获取文件夹中的图片文件并按数字排序
            if self._load_image_files():
                break
            else:
                print("❌ 文件夹中没有找到有效的图片文件")
                continue

        # 图片保存路径
        self.image_save_path = input("请输入图片保存路径: ").strip()
        if not os.path.exists(self.image_save_path):
            try:
                os.makedirs(self.image_save_path, exist_ok=True)
                print(f"✅ 已创建图片目录: {self.image_save_path}")
            except Exception as e:
                print(f"❌ 创建图片目录失败: {e}")
                return False

        # 标签保存路径
        self.label_save_path = input("请输入标签保存路径: ").strip()
        if not os.path.exists(self.label_save_path):
            try:
                os.makedirs(self.label_save_path, exist_ok=True)
                print(f"✅ 已创建标签目录: {self.label_save_path}")
            except Exception as e:
                print(f"❌ 创建标签目录失败: {e}")
                return False

        # 设置开始标注的图片编号
        while True:
            try:
                start_num = input(f"请输入从第几张图片开始标注 (1-{len(self.image_files)}): ").strip()
                self.start_image_number = int(start_num)
                if 1 <= self.start_image_number <= len(self.image_files):
                    self.current_image_index = self.start_image_number - 1  # 转换为索引
                    break
                else:
                    print(f"❌ 请输入 1 到 {len(self.image_files)} 之间的数字")
            except ValueError:
                print("❌ 请输入有效的数字")

        print(f"✅ 图片文件夹路径: {self.image_folder_path}")
        print(f"✅ 图片保存路径: {self.image_save_path}")
        print(f"✅ 标签保存路径: {self.label_save_path}")
        print(f"✅ 开始标注图片: 第{self.start_image_number}张 ({self.image_files[self.current_image_index]})")
        return True

    def _load_image_files(self):
        """加载并排序图片文件"""
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        image_files = []

        for file in os.listdir(self.image_folder_path):
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in valid_extensions:
                # 检查文件名是否为数字
                file_name = os.path.splitext(file)[0]
                try:
                    int(file_name)  # 尝试转换为数字
                    image_files.append(file)
                except ValueError:
                    print(f"⚠️  跳过非数字文件名: {file}")

        if not image_files:
            return False

        # 按数字排序
        def numeric_sort_key(filename):
            return int(os.path.splitext(filename)[0])

        image_files.sort(key=numeric_sort_key)
        self.image_files = image_files

        print(f"✅ 找到 {len(image_files)} 个图片文件")
        print("前10个文件:")
        for i, filename in enumerate(image_files[:10]):
            print(f"  {i+1}. {filename}")
        if len(image_files) > 10:
            print(f"  ... 还有 {len(image_files) - 10} 个文件")

        return True
    
    def load_current_image(self):
        """加载当前索引的图片"""
        if self.current_image_index >= len(self.image_files):
            print("✅ 所有图片已标注完成！")
            return False

        current_file = self.image_files[self.current_image_index]
        image_path = os.path.join(self.image_folder_path, current_file)

        self.current_image = cv2.imread(image_path)
        if self.current_image is not None:
            self.current_image_path = image_path
            print(f"✅ 已加载图片: {current_file} (第{self.current_image_index + 1}张)")
            print(f"图片尺寸: {self.current_image.shape[1]}x{self.current_image.shape[0]}")
            return True
        else:
            print(f"❌ 无法读取图片文件: {current_file}")
            return False

    def next_image(self):
        """切换到下一张图片"""
        self.current_image_index += 1
        return self.load_current_image()

    def previous_image(self):
        """切换到上一张图片"""
        if self.current_image_index > 0:
            self.current_image_index -= 1
            return self.load_current_image()
        else:
            print("⚠️  已经是第一张图片")
            return True
    
    def mouse_callback(self, event, x, y, flags, param):
        """鼠标回调函数"""
        if event == cv2.EVENT_LBUTTONDOWN:
            if self.current_label_id is None:
                print("⚠️  请先选择要标注的标签ID")
                return
                
            if self.temp_point is None:
                # 第一次点击，记录左上角点
                self.temp_point = (x, y)
                print(f"已选择左上角点: ({x}, {y})")
            else:
                # 第二次点击，记录右下角点并创建标注框
                x1, y1 = self.temp_point
                x2, y2 = x, y
                
                # 确保坐标正确（左上角和右下角）
                left = min(x1, x2)
                top = min(y1, y2)
                right = max(x1, x2)
                bottom = max(y1, y2)
                
                # 转换为YOLO格式 (相对坐标)
                img_h, img_w = self.current_image.shape[:2]
                center_x = (left + right) / 2.0 / img_w
                center_y = (top + bottom) / 2.0 / img_h
                width = (right - left) / img_w
                height = (bottom - top) / img_h
                
                # 添加标注框
                box_info = {
                    'label_id': self.current_label_id,
                    'center_x': center_x,
                    'center_y': center_y,
                    'width': width,
                    'height': height,
                    'pixel_coords': (left, top, right, bottom)
                }
                self.current_boxes.append(box_info)
                
                print(f"✅ 已添加标注框: 标签{self.current_label_id}({self.labels[self.current_label_id]}) "
                      f"坐标({left},{top},{right},{bottom})")
                
                self.temp_point = None
                self.draw_image()
    
    def draw_image(self):
        """绘制图片和标注框"""
        if self.current_image is None:
            return
            
        display_img = self.current_image.copy()
        
        # 绘制已有的标注框
        for box in self.current_boxes:
            left, top, right, bottom = box['pixel_coords']
            label_id = box['label_id']
            label_name = self.labels[label_id]
            
            # 绘制矩形框
            cv2.rectangle(display_img, (left, top), (right, bottom), (0, 255, 0), 2)
            
            # 绘制标签文本
            label_text = f"{label_id}: {label_name}"
            cv2.putText(display_img, label_text, (left, top-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        # 绘制临时点
        if self.temp_point is not None:
            cv2.circle(display_img, self.temp_point, 5, (0, 0, 255), -1)
            
        cv2.imshow(self.window_name, display_img)
    
    def save_labels(self):
        """保存标签文件"""
        # 获取图片文件名（不含扩展名）
        image_name = Path(self.current_image_path).stem
        label_file_path = os.path.join(self.label_save_path, f"{image_name}.txt")

        try:
            if self.current_boxes:
                # 有标注框时保存标注
                with open(label_file_path, 'w') as f:
                    for box in self.current_boxes:
                        line = f"{box['label_id']} {box['center_x']:.6f} {box['center_y']:.6f} {box['width']:.6f} {box['height']:.6f}\n"
                        f.write(line)
                print(f"✅ 已保存标签文件: {label_file_path} (包含{len(self.current_boxes)}个标注框)")
            else:
                # 没有标注框时创建空文件
                with open(label_file_path, 'w') as f:
                    pass  # 创建空文件
                print(f"✅ 已保存空标签文件: {label_file_path}")
        except Exception as e:
            print(f"❌ 保存标签文件失败: {e}")
    
    def copy_image(self):
        """复制图片到保存目录"""
        try:
            image_name = os.path.basename(self.current_image_path)
            dest_path = os.path.join(self.image_save_path, image_name)
            
            if self.current_image_path != dest_path:
                shutil.copy2(self.current_image_path, dest_path)
                print(f"✅ 已复制图片到: {dest_path}")
        except Exception as e:
            print(f"❌ 复制图片失败: {e}")
    
    def start_labeling(self):
        """开始标注过程"""
        if not self.load_current_image():
            return

        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self.mouse_callback)

        print("\n" + "=" * 60)
        print("开始标注 - 操作说明:")
        print("1. 输入标签ID开始标注该类别")
        print("2. 在图片上点击两次确定矩形框（左上角 -> 右下角）")
        print("3. 输入 'next' 保存当前标注并切换到下一张图片")
        print("4. 输入 'prev' 保存当前标注并切换到上一张图片")
        print("5. 输入 'save' 仅保存当前图片的标注")
        print("6. 输入 'skip' 跳过当前图片（不保存标注）")
        print("7. 输入 'quit' 退出程序")
        print("8. 输入 'clear' 清除当前图片的所有标注")
        print("=" * 60)

        while True:
            # 重置当前图片的标注
            self.current_boxes = []
            self.temp_point = None
            self.current_label_id = None
            self.draw_image()

            # 检查是否有已存在的标签文件
            image_name = Path(self.current_image_path).stem
            existing_label_path = os.path.join(self.label_save_path, f"{image_name}.txt")
            if os.path.exists(existing_label_path):
                print(f"⚠️  发现已存在的标签文件: {existing_label_path}")
                choice = input("是否加载已有标注? (y/n): ").strip().lower()
                if choice == 'y':
                    self._load_existing_labels(existing_label_path)
                    self.draw_image()

            while True:
                print(f"\n当前图片: {os.path.basename(self.current_image_path)} "
                      f"(第{self.current_image_index + 1}/{len(self.image_files)}张)")
                print(f"当前标注框数量: {len(self.current_boxes)}")
                print(f"当前标签列表:")
                for label_id, label_name in sorted(self.labels.items()):
                    print(f"  {label_id}: {label_name}")

                command = input("\n请输入命令 (标签ID/next/prev/save/skip/quit/clear): ").strip()

                if command.lower() == 'quit':
                    cv2.destroyAllWindows()
                    return
                elif command.lower() == 'next':
                    self.save_labels()
                    self.copy_image()
                    if self.next_image():
                        break  # 跳出内层循环，加载下一张图片
                    else:
                        cv2.destroyAllWindows()
                        return  # 所有图片处理完成
                elif command.lower() == 'prev':
                    self.save_labels()
                    self.copy_image()
                    if self.previous_image():
                        break  # 跳出内层循环，加载上一张图片
                elif command.lower() == 'save':
                    self.save_labels()
                    self.copy_image()
                elif command.lower() == 'skip':
                    if self.next_image():
                        break  # 跳过当前图片，加载下一张
                    else:
                        cv2.destroyAllWindows()
                        return
                elif command.lower() == 'clear':
                    self.current_boxes = []
                    self.temp_point = None
                    self.current_label_id = None
                    self.draw_image()
                    print("✅ 已清除所有标注")
                else:
                    try:
                        label_id = int(command)
                        if label_id in self.labels:
                            self.current_label_id = label_id
                            self.temp_point = None
                            print(f"✅ 当前标注标签: {label_id} ({self.labels[label_id]})")
                            print("请在图片上点击两次确定矩形框")
                        else:
                            print(f"❌ 标签ID {label_id} 不存在")
                    except ValueError:
                        print("❌ 无效命令")

                # 处理OpenCV窗口事件
                key = cv2.waitKey(1) & 0xFF
                if key == 27:  # ESC键退出
                    cv2.destroyAllWindows()
                    return

        cv2.destroyAllWindows()

    def _load_existing_labels(self, label_file_path):
        """加载已存在的标签文件"""
        try:
            with open(label_file_path, 'r') as f:
                lines = f.readlines()

            self.current_boxes = []
            img_h, img_w = self.current_image.shape[:2]

            for line in lines:
                line = line.strip()
                if not line:
                    continue

                parts = line.split()
                if len(parts) == 5:
                    label_id = int(parts[0])
                    center_x = float(parts[1])
                    center_y = float(parts[2])
                    width = float(parts[3])
                    height = float(parts[4])

                    # 转换为像素坐标
                    left = int((center_x - width/2) * img_w)
                    top = int((center_y - height/2) * img_h)
                    right = int((center_x + width/2) * img_w)
                    bottom = int((center_y + height/2) * img_h)

                    box_info = {
                        'label_id': label_id,
                        'center_x': center_x,
                        'center_y': center_y,
                        'width': width,
                        'height': height,
                        'pixel_coords': (left, top, right, bottom)
                    }
                    self.current_boxes.append(box_info)

            print(f"✅ 已加载 {len(self.current_boxes)} 个标注框")
        except Exception as e:
            print(f"❌ 加载标签文件失败: {e}")
    
    def run(self):
        """运行标注工具"""
        print("=" * 60)
        print("YOLOv5 标注工具")
        print("=" * 60)
        
        # 设置标签
        if not self.setup_labels():
            return
            
        # 设置路径
        if not self.setup_paths():
            return
        
        # 开始标注
        self.start_labeling()
        
        print("\n✅ 标注工具已退出")

def main():
    """主函数"""
    tool = YOLOv5LabelingTool()
    tool.run()

if __name__ == "__main__":
    main()
